import os
import cv2
import tensorflow as tf
import tqdm as tqdm

if __name__ == '__main__':
    # 需要处理的数据集根目录
    Dataset_path = 'Split_Dataset'
    # 用于生成需要写入的文件路径
    file_names = [i.split('.')[0] for i in os.listdir(os.path.join(Dataset_path, 'DSM'))]
    dsm_filename = [os.path.join(Dataset_path, 'DSM', i + '.tif') for i in file_names]
    label_filename = [os.path.join(Dataset_path, 'Label', i + '.png') for i in file_names]
    rgb_filename = [os.path.join(Dataset_path, 'RGB', i + '.png') for i in file_names]
    # 开始进行TFRecords文件的构建
    with tf.io.TFRecordWriter('Potsdam.tfrecords') as writer:
        tqdm_file = tqdm.tqdm(iterable=zip(dsm_filename, rgb_filename, label_filename), total=len(dsm_filename))
        for dsm, rgb, label in tqdm_file:
            # 进行高程图的读取
            dsm_data = cv2.imread(dsm, -1)
            dsm_data = dsm_data.ravel()

            # 进行光学图的读取
            rgb_data = open(rgb, 'rb').read()

            # 进行标签的读取
            label_data = open(label, 'rb').read()

            # 建立tf.train.Feature字典
            feature = {
                'dsm': tf.train.Feature(float_list=tf.train.FloatList(value=dsm_data)),
                'rgb': tf.train.Feature(bytes_list=tf.train.BytesList(value=[rgb_data])),
                'label': tf.train.Feature(bytes_list=tf.train.BytesList(value=[label_data]))
            }

            # 通过字典创建Example
            example = tf.train.Example(features=tf.train.Features(feature=feature))
            # 将Example序列化并写入TFRecords文件
            writer.write(example.SerializeToString())
        tqdm_file.close()
